# -*- coding: utf-8 -*-
# @Author: Cursor 
# @Date: 2025-01-22
# @Last Modified by: Tim Liu
# @Last Modified time: 2025-01-22

# @Description: MongoDB Conversation Service, manages the conversation between the user and the bot using MongoDB.

import uuid
from datetime import datetime

from langchain_mongodb import MongoDBChatMessageHistory
from crewplus.apps.conversation.schemas import Conversation
from .conversation_service import ConversationService
from typing import List
from config.settings import *

DEFAULT_COLLECTION = "chat_histories"

class MongoDBConversationService(ConversationService):
    """
    MongoDB Conversation Service, manages the conversation between the user and the bot (session separation) using MongoDB.
    """

    def __init__(self, model: str = None, connection_string: str = MONGO_DB_CONN):
        super().__init__(model)
        self.connection_string = connection_string
        self.histories = {}

    def get_history(self, session_id: str) -> MongoDBChatMessageHistory:
        if session_id not in self.histories:
            self.histories[session_id] = MongoDBChatMessageHistory(
                session_id=session_id,
                connection_string=self.connection_string,
                database_name=MONGO_DB_NAME,
                collection_name=DEFAULT_COLLECTION,
            )

        return self.histories[session_id]

    def new_conversation(self, user_id: int = 0):
        # Create a new conversation in the database, save the conversation_id, user_id and start time
        data = Conversation()
        data.session_id = uuid.uuid4()
        data.create_user_id = user_id
        data.converse_start_time = str(datetime.now())

        return data

    def add_message(self, conversation_id: str, message: str = None, type: str = "user"):
        if message is None:
            return

        history = self.get_history(conversation_id)
        
        if type == "user":
            history.add_user_message(message)
        else:
            history.add_ai_message(message)
        
        return history

    def add_messages(self, conversation_id: str, messages: List[dict]):
        history = self.get_history(conversation_id)
        history.add_messages(messages)

    def add_user_message(self, conversation_id: str, human_message: str):
        history = self.get_history(conversation_id)
        history.add_user_message(human_message)

    def add_ai_message(self, conversation_id: str, ai_message: str):
        history = self.get_history(conversation_id)
        history.add_ai_message(ai_message)

    def get_history_messages(self, conversation_id: str):
        history = self.get_history(conversation_id)
        return history.messages